# GRU是LSTM的一种变形版本，它将遗忘门和输入门组合成一个更新门，并且合并了单元状态和隐藏状态，模型更加简单。
import torch
from torch import lstm
import torch.nn as nn

batch_size = 10
seq_len = 20  # 句子长度
vocab_size = 100  # 字典长度
embedding_dim = 30  # 表示词语的向量长度
hidden_size = 18
num_layer = 1


def test1():
    # input_size，设置输入数据的形状，即embedding_dim，这里表示告诉函数embedding处理过的数据形状格式，这样LSTM才能正确读取数据。
    # hidden_size，隐藏层的数量，即每一层有多少个LSTM单元，LSTM这个对象里包含了一组LSTM单元。
    # num_layer，RNN中LSTM单元的层数。
    # batch_first，把batch_size放了开头。
    # dropout，默认0，设置部分参数随机失活，降低过拟合。
    # bidirectional，是否使用双向LSTM
    # LSTM对象传入数据的包括当前数据，前一次的隐藏状态h_n，前一次的存储c_n。

    dataset = torch.randint(low=0, high=100, size=[batch_size, seq_len])
    print(dataset.size())

    embedding = nn.Embedding(vocab_size, embedding_dim=embedding_dim)
    input_embedded = embedding(dataset)
    print(input_embedded.size())

    lstm = nn.GRU(input_size=embedding_dim, hidden_size=hidden_size, num_layers=num_layer, batch_first=True,
                  bidirectional=False)
    output, h_n = lstm(input_embedded)
    print(output.size())
    print(h_n.size())

    # 最后一个时间步的输出，output中间那个维度的最后一个值
    last_output = output[:, -1, :]
    # 最后一次hidden_state
    last_hidden_state = h_n[-1, :, :]
    print(last_output)
    print(last_hidden_state)
    print(last_output == last_hidden_state)  # 其实这两个是相同的


if __name__ == '__main__':
    test1()
